#!/usr/bin/python
# -*- coding: utf-8 -*-

from pandas import DataFrame
from numpy import nan

from core.dataClasses import StockTradeDataColumnName
from core.index_base import IndexBase

"""
思想：在股票市场上，买方和卖方力量的消长会影响股票的价格。如果股票的卖方力量大于卖方力量，则股票的价格会上涨；如果股票的卖方力量大于买方力量，则股票的价格会下跌。
RSI（Relative Strength Index）相对强弱指标，是用一种特定公式计算出来的值，投资者可以通过RSI的取值来判断股票的买入和卖出情况，进而预测未来股票的价格走势。
    RSI = 100 - 100/(1 + RS) 或 RSI = 100 * RS/(1 + RS)
    RS = UP/DOWN
    将RS = UP/DOWN带入RSI = 100 * RS/(1 + RS)可以推出：
        RSI = 100 * (UP/DOWN)/(1 + UP/DOWN) = 100 * UP/(UP + DOWN)
    其中：
        RSI表示相对强弱指标值；若t表示所参考数据的期数，UP表示t期内估价上涨幅度的平均值，DOWN表示t期内估价下跌幅度的平均值。
        在计算UP和DOWN的值时，需要求出上涨幅度的平均值和下跌幅度的平均值，一般使用算数平均值来求UP和DOWN。
            均值可分为简单移动平均值（SMA）、加权移动平均值（WMA）和指数移动平均值（EMA）。
策略：RSI的取值大于50越多，则表明股票上涨的力量超过下跌的力量更大。当RSI的取值小于50越多，则股票的下跌力量超过上涨力量更大。
RSI天数的差异：
    韦尔斯 维尔德（Wells Wilder）指出，通过运用月周期28日的一般来计算RSI的值进行预测是有效的，他推荐RSI的时间跨度默认值为14日。
    一些看盘软件设有6日RSI、12日RSI和24日RSI指标。6日近似一周的时间周期，12日可以看做半个月的时间跨度，而24日约为一个月的时间跨度。
超买线、超卖线、中心线
    RSI取值等于80或20分别为较常见的“超买线”和“超卖线”的刻画。RSI为80是股票超买的临界点，RSI为20是股票超卖的临界点，RSI取值为50设定为“中心线“，该线表明股票的买入力量等于卖出力量。
    当RSI大于80时，股票出现超买信号。股票买入力量过大，买入力量在未来可能会减小，所以股票未来价格可能会下跌，此时卖出股票，未来下跌后再买入股票，从而赚取价差。
    当RSI小于20时，股票出现超卖信号。股票卖出力量过大，卖出力量在未来可能会回归正常，所以股票未来价格可能会上涨，投资者此时可以做多股票，未来价格上涨后再卖出。
黄金交叉和死亡交叉：
    当短期RSI线向上穿过长期RSI线时，股票近期买入的力量较强，价格上涨的力量很大，其释放出一个较强的买入信号，这个信号被成为“黄金交叉”。
    当短期RSI线向下跌破长期RSI线时，股票近期卖出的力量较强，价格下跌的力量很大，其释放出一个较强的卖出信号，被成为“死亡交叉”。
"""
SMA_RSI_COLUMN_PREFIX = 'sma_rsi'


class SmaRSI(IndexBase):
    def __init__(self, period=6, key=StockTradeDataColumnName.CLOSE):
        super(SmaRSI, self).__init__()
        self.period = period
        self.key = key

    def compute(self, data: DataFrame):
        column = SMA_RSI_COLUMN_PREFIX + str(self.period)
        self.added_keys.append(column)

        data[column] = nan
        data['change+'] = data[data[StockTradeDataColumnName.CHANGE] > 0][StockTradeDataColumnName.CHANGE]
        data['change+'] = data['change+'].replace(to_replace=nan, value=0)
        data['change-'] = data[data[StockTradeDataColumnName.CHANGE] < 0][StockTradeDataColumnName.CHANGE]
        data['change-'] = data['change-'].replace(to_replace=nan, value=0)
        data['up'] = nan
        data['up'] = data['change+'].rolling(self.period).mean()
        data['down'] = nan
        data['down'] = data['change-'].abs().rolling(self.period).mean()
        data[column] = 100 * data['up']/(data['up'] + data['down'])
        data = data.drop(labels=['change+', 'change-', 'up', 'down'], axis=1)
        # print(data)


if __name__ == "__main__":
    from domain.transaction_data.repository import transaction_data_repository

    SmaRSI(period=20).compute(transaction_data_repository.get_contain_latest_data_by_tushare('002044.SZ'))